package algorithm

import org.apache.spark.graphx.{Graph, VertexId, VertexRDD}

object ShortestPathAlgorithm {
  //Note 新生成的图 边变量ED是其权重 顶点变量VD是一个二元的元组
  //其中_1代表是否遍历过 _2代表当前顶点的距离
  def dijkstra[VD](g: Graph[VD, Double], origin: VertexId, direct: Boolean): Graph[(VD, Double), Double] ={
    //初始化 对顶点进行变换 变换后顶点属性VD变成一个二元元组 (bool,double)
    var g2 = g.mapVertices {
      case (vid, _) =>
        val vd:Double = if (vid == origin) 0 else Double.MaxValue
        (false, vd)
    }

    //遍历所有的点
    (0L until g.vertices.count()).foreach{
      i:Long =>{
        //定位未遍历节点中 距离最小的
        val currentVertexId: VertexId = g2.vertices.filter(!_._2._1)
          .fold((0L, (false, Double.MaxValue))) {
            case (a, b) => if (a._2._2 < b._2._2) a else b
          }._1

        //有当前节点往相邻节点发消息
        val newDistances :VertexRDD[Double] = g2.aggregateMessages[Double](
          //sendMes 向邻边发送消息 内容为边的距离与最短路径之和
          ctx=>{
            if(ctx.srcId ==currentVertexId) ctx.sendToDst(ctx.srcAttr._2 +ctx.attr)
            if(ctx.dstId ==currentVertexId && !direct) ctx.sendToSrc(ctx.dstAttr._2 +ctx.attr) //考虑方向性
          },
          //mergeMsg 选择较小的值作为当前顶点的最短路径值
          (a,b)=>math.min(a,b))

        //生成结果图
        g2 =g2.outerJoinVertices(newDistances){
          (vid, vd, newSum)=>{
            (vd._1 || vid == currentVertexId,math.min(vd._2,newSum.getOrElse(Double.MaxValue)))
          }
        }
      }
    }

    g.outerJoinVertices(g2.vertices){
      (_, vd, dist)=>{
        (vd,dist.getOrElse((false,Double.MaxValue))._2)
      }
    }
  }

  //Note 新生成的图 边变量ED是其权重 顶点变量VD是一个三元的元组
  //其中_1代表是否遍历过 _2代表当前顶点的距离 _3 代表遍历路径
  def dijkstraWithTrace[VD](g: Graph[VD, Double], origin: VertexId, direct: Boolean): Graph[(VD, Double,List[VertexId]), Double] = {
    //初始化 对顶点进行变换 变换后顶点属性VD变成一个三元元组 (bool,double,list)

    var g2 = g.mapVertices {
      case (vid, _) =>
        val vd:Double = if (vid == origin) 0 else Double.MaxValue
        (false, vd,List[VertexId]())
    }

    //遍历所有的点
    (0L until g.vertices.count()).foreach{
      i:Long =>{
        //定位未遍历节点中 距离最小的
        val currentVertexId: VertexId = g2.vertices.filter(!_._2._1)
          .fold((0L, (false, Double.MaxValue,List[VertexId]()))) {
            case (a, b) => if (a._2._2 < b._2._2) a else b
          }._1

        //有当前节点往相邻节点发消息
        val newDistances :VertexRDD[(Double,List[VertexId])] = g2.aggregateMessages[(Double,List[VertexId])](
          //sendMes 向邻边发送消息 内容为边的距离与最短路径之和
          ctx=>{
            if(ctx.srcId ==currentVertexId) ctx.sendToDst((ctx.srcAttr._2 +ctx.attr,ctx.srcAttr._3:+ctx.srcId))
            if(ctx.dstId ==currentVertexId && !direct) ctx.sendToSrc(ctx.dstAttr._2 +ctx.attr,ctx.dstAttr._3:+ctx.dstId) //考虑方向性
          },
          //mergeMsg 选择较小的值作为当前顶点的最短路径值
          (a,b)=>if (a._1<b._1) a else b
        )

        //生成结果图
        g2 =g2.outerJoinVertices(newDistances){
          (vid, vd, newSum)=>{
            val newSumVal = newSum.getOrElse((Double.MaxValue,List[VertexId]()))
            (vd._1 || vid == currentVertexId,math.min(vd._2,newSumVal._1),if (vd._2<newSumVal._1) vd._3 else newSumVal._2)
          }
        }
      }
    }

    g.outerJoinVertices(g2.vertices){
      (_, vd, dist)=>{
        val newDistVal = dist.getOrElse((false,Double.MaxValue,List[VertexId]()))
        (vd,newDistVal._2,newDistVal._3)
      }
    }
  }

  //Note 新生成的图 边变量ED是其权重 顶点变量VD是一个二元的元组
  //其中_1代表是否遍历过 _2代表当前顶点的距离
  def pregel[VD](g: Graph[VD, Double], origin: VertexId,direct: Boolean,maxIterations:Integer)
  : Graph[(VD, Double), Double] = {
    //初始化遍历前的图
    val graph = g.mapVertices((vid, _) =>{
      val vd:Double = if (vid == origin) 0 else Double.MaxValue //除根节点外 其余的最大距离为最大值
      vd
    })

    // Initialize vertices with an initial distance of Infinity
    val bfs = graph.pregel(Double.MaxValue, maxIterations)(
      // vprog When a vertex receives a message, it checks the candidate distance from root against the known best route
      //每次迭代 vprog 其次执行 用于处理本轮收到的最短距离消息
      (vid, vAttr, vMsg) => {
        if (vMsg < vAttr) vMsg else vAttr
      },
      //sendMsg
      triplet => {
        //遍历所有的edge根据src顶点的最短距离计算dst顶点的最短距离 并发消息
        if (triplet.srcAttr != Double.MaxValue &&
          triplet.srcAttr + triplet.attr < triplet.dstAttr) {
          Iterator((triplet.dstId, triplet.srcAttr + triplet.attr))
        } else if (triplet.dstAttr != Double.MaxValue &&
          triplet.dstAttr + triplet.attr < triplet.srcAttr && !direct){
          Iterator((triplet.srcId, triplet.dstAttr + triplet.attr)) //针对无方向的 可以向src节点发送消息
        }
        else {
          Iterator.empty
        }
      },
      //mergeMsg
      //Because a vertex might receive multiple messages in a single level, we want to reduce this to just the shortest path
      //每一迭代 mergeMsg 首先执行 由于节点会收到多个节点发送来的最短距离 我们通过reduce获取最短的
      (a, b) => if (a < b) a else b
    )

    g.outerJoinVertices(bfs.vertices){
      (_, vd, dist)=>{
        (vd,dist.getOrElse(Double.MaxValue)) //vd 原有的顶点属性 dist 根据图计算后的距离
      }
    }
  }

  //Note 新生成的图 边变量ED是其权重 顶点变量VD是一个三元的元组
  //其中_1代表是否遍历过 _2代表当前顶点的距离 _3 代表遍历路径
  def pregelWithTrace[VD](g: Graph[VD, Double], origin: VertexId,direct: Boolean,maxIterations:Integer)
  : Graph[(VD, Double,List[VertexId]), Double] = {
    //初始化遍历前的图
    val graph = g.mapVertices((vid, _) =>{
      val vd:Double = if (vid == origin) 0 else Double.MaxValue
      (vd,List[VertexId]()) //除根节点外 其余的最大距离为最大值 遍历列表为空
    })

    // Initialize vertices with an initial distance of Infinity
    val bfs = graph.pregel((Double.MaxValue,List[VertexId]()), maxIterations)(
      // vprog When a vertex receives a message, it checks the candidate distance from root against the known best route
      //每次迭代 vprog 其次执行 用于处理本轮收到的最短距离消息
      (vid, vAttr, vMsg) => {
        if (vMsg._1 < vAttr._1) vMsg else vAttr
      },
      //sendMsg
      triplet => {
        //遍历所有的edge根据src顶点的最短距离计算dst顶点的最短距离 并发消息
        if (triplet.srcAttr._1 != Double.MaxValue &&
          triplet.srcAttr._1 + triplet.attr < triplet.dstAttr._1) {
          Iterator((triplet.dstId, (triplet.srcAttr._1 + triplet.attr,triplet.srcAttr._2:+triplet.srcId)))
        } else if (triplet.dstAttr._1 != Double.MaxValue &&
          triplet.dstAttr._1 + triplet.attr < triplet.srcAttr._1 && !direct){
          Iterator((triplet.srcId, (triplet.dstAttr._1 + triplet.attr,triplet.dstAttr._2:+triplet.dstId))) //针对无方向的 可以向src节点发送消息
        }
        else {
          Iterator.empty
        }
      },
      //mergeMsg
      //Because a vertex might receive multiple messages in a single level, we want to reduce this to just the shortest path
      //每一迭代 mergeMsg 首先执行 由于节点会收到多个节点发送来的最短距离 我们通过reduce获取最短的
      (a, b) => if (a._1 < b._1) a else b
    )

    g.outerJoinVertices(bfs.vertices){
      (_, vd, dist)=>{
        val newDistVal = dist.getOrElse((Double.MaxValue,List[VertexId]()))
        (vd,newDistVal._1,newDistVal._2)
      }
    }
  }
}
